#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
重建MongoDB全文搜索索引
优化中文搜索支持
"""

import json
import os
from datetime import datetime
from pymongo import MongoClient, TEXT, ASCENDING, DESCENDING
from typing import List, Dict, Any
import jieba

class SearchIndexOptimizer:
    """搜索索引优化器"""
    
    def __init__(self, connection_string: str = "mongodb://localhost:27017/", db_name: str = "civil_code_db"):
        self.connection_string = connection_string
        self.db_name = db_name
        
    def rebuild_text_indexes(self) -> bool:
        """重建文本搜索索引"""
        print("🔄 开始重建全文搜索索引...")
        
        try:
            client = MongoClient(self.connection_string)
            db = client[self.db_name]
            collection = db['civil_code_articles']
            
            # 1. 删除现有的文本索引
            print("🗑️ 删除现有文本索引...")
            existing_indexes = list(collection.list_indexes())
            
            for index in existing_indexes:
                index_name = index.get('name')
                if index_name and 'text' in index_name.lower():
                    try:
                        collection.drop_index(index_name)
                        print(f"   ✅ 删除索引: {index_name}")
                    except Exception as e:
                        print(f"   ⚠️ 删除索引失败 {index_name}: {str(e)}")
            
            # 2. 创建优化的文本索引
            print("🔧 创建优化的文本搜索索引...")
            
            # 配置文本索引字段和权重
            text_index_spec = [
                ("content", "text"),      # 条文内容 - 权重10
                ("title", "text"),        # 条文标题 - 权重5  
                ("explanation", "text"),  # 条文解释 - 权重3
                ("keywords", "text")      # 关键词 - 权重8
            ]
            
            # 创建文本索引，设置语言为none以支持中文
            collection.create_index(
                text_index_spec,
                name="civil_code_text_search",
                default_language='none',  # 支持多语言，特别是中文
                weights={
                    "content": 10,
                    "keywords": 8,
                    "title": 5,
                    "explanation": 3
                }
            )
            
            print("   ✅ 文本搜索索引创建成功")
            
            # 3. 创建关键词数组索引
            print("🔧 创建关键词数组索引...")
            collection.create_index(
                [("keywords", ASCENDING)],
                name="keywords_array_index"
            )
            print("   ✅ 关键词索引创建成功")
            
            # 4. 测试文本搜索功能
            print("🧪 测试文本搜索功能...")
            self._test_text_search(collection)
            
            client.close()
            return True
            
        except Exception as e:
            print(f"❌ 索引重建失败: {str(e)}")
            return False
    
    def create_performance_indexes(self) -> bool:
        """创建性能优化索引"""
        print("🚀 创建性能优化索引...")
        
        try:
            client = MongoClient(self.connection_string)
            db = client[self.db_name]
            collection = db['civil_code_articles']
            
            # 1. 条文号唯一索引
            print("🔧 创建条文号唯一索引...")
            collection.create_index(
                [("article_number", ASCENDING)],
                name="article_number_unique",
                unique=True
            )
            print("   ✅ 条文号索引创建成功")
            
            # 2. 章节复合索引
            print("🔧 创建章节复合索引...")
            collection.create_index(
                [
                    ("chapter_info.part", ASCENDING),
                    ("chapter_info.chapter_number", ASCENDING)
                ],
                name="chapter_compound_index"
            )
            print("   ✅ 章节索引创建成功")
            
            # 3. 时间索引
            print("🔧 创建时间索引...")
            collection.create_index(
                [("created_at", DESCENDING)],
                name="created_at_index"
            )
            collection.create_index(
                [("updated_at", DESCENDING)],
                name="updated_at_index"
            )
            print("   ✅ 时间索引创建成功")
            
            # 4. 显示所有索引
            print("\n📋 当前所有索引:")
            indexes = list(collection.list_indexes())
            for index in indexes:
                index_name = index.get('name', 'unnamed')
                index_keys = list(index.get('key', {}).keys())
                index_unique = index.get('unique', False)
                unique_flag = " [UNIQUE]" if index_unique else ""
                print(f"   - {index_name}: {', '.join(index_keys)}{unique_flag}")
            
            client.close()
            return True
            
        except Exception as e:
            print(f"❌ 性能索引创建失败: {str(e)}")
            return False
    
    def enhance_keywords_field(self) -> bool:
        """增强关键词字段"""
        print("🔧 增强关键词字段...")
        
        try:
            client = MongoClient(self.connection_string)
            db = client[self.db_name]
            collection = db['civil_code_articles']
            
            # 获取所有文档
            documents = list(collection.find({}))
            print(f"📄 处理 {len(documents)} 条法条的关键词...")
            
            updated_count = 0
            
            for doc in documents:
                content = doc.get('content', '')
                existing_keywords = doc.get('keywords', [])
                
                # 使用jieba重新提取关键词
                enhanced_keywords = self._extract_enhanced_keywords(content)
                
                # 合并现有关键词和新提取的关键词
                all_keywords = list(set(existing_keywords + enhanced_keywords))
                
                # 更新文档
                collection.update_one(
                    {'_id': doc['_id']},
                    {'$set': {'keywords': all_keywords[:10]}}  # 最多保留10个关键词
                )
                
                updated_count += 1
                
                if updated_count % 100 == 0:
                    print(f"   📊 已处理 {updated_count}/{len(documents)} 条")
            
            print(f"✅ 关键词字段增强完成，更新了 {updated_count} 条法条")
            
            client.close()
            return True
            
        except Exception as e:
            print(f"❌ 关键词字段增强失败: {str(e)}")
            return False
    
    def _extract_enhanced_keywords(self, content: str) -> List[str]:
        """提取增强的关键词"""
        # 法律专业词汇
        legal_terms = {
            '民事主体', '自然人', '法人', '非法人组织', '民事权利', '民事义务',
            '民事责任', '诉讼时效', '代理', '监护', '宣告失踪', '宣告死亡',
            '个体工商户', '农村承包经营户', '营利法人', '非营利法人',
            '物权', '所有权', '用益物权', '担保物权', '占有', '共有',
            '建设用地使用权', '宅基地使用权', '居住权', '地役权',
            '抵押权', '质权', '留置权', '合同', '要约', '承诺',
            '违约责任', '合同解除', '买卖合同', '赠与合同', '借贷合同',
            '租赁合同', '承揽合同', '委托合同', '保证合同',
            '人格权', '生命权', '身体权', '健康权', '姓名权', '肖像权',
            '名誉权', '荣誉权', '隐私权', '婚姻', '家庭', '收养',
            '继承', '遗嘱', '遗赠', '侵权责任', '过错责任', '无过错责任',
            '合法权益', '法律保护', '人身关系', '财产关系'
        }
        
        # 使用jieba分词
        words = jieba.lcut(content)
        
        # 提取关键词
        keywords = []
        
        # 优先提取法律专业词汇
        for word in words:
            if word in legal_terms and word not in keywords:
                keywords.append(word)
        
        # 提取其他有意义的词汇
        for word in words:
            if (len(word) >= 2 and 
                word not in keywords and 
                word not in ['规定', '条款', '法律', '本条', '应当', '可以', '依照', '按照', '根据']):
                keywords.append(word)
                if len(keywords) >= 8:
                    break
        
        return keywords
    
    def _test_text_search(self, collection) -> None:
        """测试文本搜索功能"""
        test_queries = ['民事主体', '合法权益', '人身关系', '财产权利']
        
        for query in test_queries:
            try:
                results = list(collection.find(
                    {"$text": {"$search": query}},
                    {"article_number": 1, "title": 1, "score": {"$meta": "textScore"}}
                ).sort([("score", {"$meta": "textScore"})]).limit(3))
                
                if results:
                    print(f"   ✅ 搜索 '{query}': 找到 {len(results)} 条相关法条")
                    for result in results:
                        score = result.get('score', 0)
                        print(f"      - {result['article_number']} (得分: {score:.2f})")
                else:
                    print(f"   ⚠️ 搜索 '{query}': 无结果")
            except Exception as e:
                print(f"   ❌ 搜索 '{query}' 失败: {str(e)}")

def main():
    """主函数"""
    print("=" * 60)
    print("🔄 重建MongoDB全文搜索索引")
    print("=" * 60)
    
    optimizer = SearchIndexOptimizer()
    
    success = True
    
    # 1. 重建文本索引
    if not optimizer.rebuild_text_indexes():
        success = False
    
    # 2. 创建性能索引
    if not optimizer.create_performance_indexes():
        success = False
    
    # 3. 增强关键词字段
    if not optimizer.enhance_keywords_field():
        success = False
    
    if success:
        print("\n🎉 全文搜索索引重建完成！")
    else:
        print("\n❌ 索引重建过程中出现错误")

if __name__ == "__main__":
    main()
